import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
from sklearn.linear_model import LogisticRegression


# 特征工程，给字段值赋予权重
def data_processing(pth):
    df = pd.read_csv(pth)
    df.drop(['Over18', 'StandardHours', 'EmployeeNumber', 'Gender'], axis='columns', inplace=True)
    # 特征工程，给离散值赋予权重
    # 然后用ration值替换离散值
    cols = df.columns.drop('Attrition')
    for item in cols.values:
        df1 = pd.crosstab(df[item], df['Attrition']).reset_index()
        df1['ratio'] = df1[1] / (df1[1] + df1[0])
        df = pd.merge(df, df1[[item, 'ratio']], on=item)
        df.drop(item, axis='columns', inplace=True)
        df.rename(columns={'ratio': item}, inplace=True)
    return df

# 进行模型训练
def train_model(data_processed):
    df = data_processed.copy()
    x = df.drop('Attrition', axis='columns')
    y = df['Attrition']
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=25)
    es = LogisticRegression(random_state=25)
    es.fit(x_train, y_train)
    y_pre = es.predict(x)
    print("当前模型的roc评分为：{}".format(roc_auc_score(y, y_pre)))
    return es


if __name__ == '__main__':
    df = data_processing('../data/train.csv')
    es = train_model(df)
